from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from kokoro import KModel, KPipeline
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
import io
import json
from typing import Optional

# 初始化FastAPI应用
app = FastAPI(title="Kokoro TTS 流式语音服务")

# 允许跨域请求（方便前端调用）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境需指定具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 模型配置
REPO_ID = 'hexgrad/Kokoro-82M-v1.1-zh'
SAMPLE_RATE = 24000  # 采样率固定为24000Hz
N_ZEROS = 5000  # 段落间静音长度（约0.2秒）
device = 'cuda' if torch.cuda.is_available() else 'cpu'

en_pipeline = KPipeline(lang_code='a', repo_id=REPO_ID, model=False)
def en_callable(text):
    if text == 'Kokoro':
        return 'kˈOkəɹO'
    elif text == 'Sol':
        return 'sˈOl'
    return next(en_pipeline(text)).phonemes

# 加载模型（全局单例，避免重复加载）
class TTSModel:
    def __init__(self):
        self.model = None
        self.zh_pipeline = None
        self.voices = {
            "zf_001": "./voices/zf_001.pt",  # 女性音色
            "zm_010": "./voices/zm_010.pt"   # 男性音色（若有）
        }
        
        self.loaded = False

    def load(self, model_path="./kokoro-v1_1-zh.pth", config_path="./config.json"):
        """加载模型（首次调用时自动加载）"""
    #     def en_callable(text):
    #         en_pipeline = KPipeline(lang_code='a', repo_id=REPO_ID, model=self.model)
    #         if text == 'Kokoro':
    #             return 'kˈOkəɹO'
    #         elif text == 'Sol':
    #             return 'sˈOl'
    #         return {"phonemes":next(en_pipeline(text)).phonemes,
    #     "voice": "./voices/af_heart.pt" ,  # 默认英文音色
    # }
        if not self.loaded:
            try:
                self.model = KModel(
                    config=config_path,
                    repo_id=REPO_ID,
                    model=model_path
                ).to(device).eval()
                self.zh_pipeline = KPipeline(lang_code='z', repo_id=REPO_ID, model=self.model,en_callable=en_callable)
                self.loaded = True
                print(f"模型加载成功，使用设备：{device}")
            except Exception as e:
                print(f"模型加载失败：{e}")
                raise HTTPException(status_code=500, detail=f"模型初始化失败：{str(e)}")

# 实例化模型（全局唯一）
tts_model = TTSModel()

# 请求参数模型
class TTSRequest(BaseModel):
    text: str  # 需要合成的文本
    voice: Optional[str] = "zf_001"  # 音色选择
    speed: Optional[float] = 1.0  # 语速（0.5-2.0）

# 语速调整函数（根据文本长度动态调整）
def speed_callable(len_ps, base_speed=1.0):
    """
    动态调整语速：文本越长，语速越慢（避免合成过于急促）
    :param len_ps: 文本长度
    :param base_speed: 基础语速（用户指定）
    """
    speed = 0.8
    if len_ps <= 83:
        speed = 1.0
    elif len_ps < 183:
        speed = 1.0 - (len_ps - 83) / 500
    return speed * base_speed * 1.1  # 保持与原代码一致的系数

# 流式生成音频的生成器函数
def tts_stream_generator(text: str, voice: str, speed: float):
    """
    流式生成音频片段，逐句返回WAV格式数据
    """
    # 确保模型已加载
    tts_model.load()
    
    # 检查音色是否存在
    if voice not in tts_model.voices:
        raise HTTPException(status_code=400, detail=f"不支持的音色：{voice}，可选音色：{list(tts_model.voices.keys())}")
    voice_path = tts_model.voices[voice]
    
    # 简单分句（按中文标点分割）
    sentences = [s.strip() for s in text.split(r'[，。,;；？！?！]') if s.strip()]
    if not sentences:
        raise HTTPException(status_code=400, detail="输入文本为空或无法分割")
    
    # 逐句生成音频
    first_sentence = True
    
    for sentence in sentences:
        try:
            # 生成当前句的音频
            generator = tts_model.zh_pipeline(
                sentence,
                voice=voice_path,
                speed=lambda len_ps: speed_callable(len_ps, base_speed=speed)

            )
            result = next(generator)
            wav = result.audio  # 音频数据（numpy数组）
            
            # 非首句前添加静音
            if not first_sentence and N_ZEROS > 0:
                silence = np.zeros(N_ZEROS, dtype=np.float32)
                wav = np.concatenate([silence, wav])
            
            first_sentence = False
            
            # 将numpy数组转为WAV格式的字节流
            buffer = io.BytesIO()
            sf.write(buffer, wav, SAMPLE_RATE, format="WAV", subtype="PCM_16")
            buffer.seek(0)
            yield buffer.read()  # 流式返回当前片段
            
        except Exception as e:
            print(f"生成音频失败（句子：{sentence}）：{e}")
            raise HTTPException(status_code=500, detail=f"生成音频失败：{str(e)}")

# TTS流式接口
@app.post("/api/tts/stream", summary="流式语音合成接口")
async def tts_stream(request: TTSRequest):
    """
    流式返回语音合成结果（WAV格式）
    - 支持逐句生成并返回音频，前端可边接收边播放
    - 音频格式：PCM_16，采样率24000Hz，单声道
    """
    try:
        # 返回流式响应，MIME类型为音频流
        return StreamingResponse(
            tts_stream_generator(request.text, request.voice, request.speed),
            media_type="audio/wav",
            headers={
                "Content-Disposition": 'attachment; filename="tts_stream.wav"',
                "Transfer-Encoding": "chunked"  # 分块传输
            }
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"服务错误：{str(e)}")

# 首页（提供简单的测试界面）
@app.get("/", response_class=HTMLResponse)
async def index():
    return """
    <html>
        <head>
            <title>Kokoro TTS 流式服务</title>
            <style>
                body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
                .container { text-align: center; }
                textarea { width: 100%; height: 150px; margin: 20px 0; padding: 10px; }
                .controls { margin: 10px 0; }
                button { padding: 10px 20px; background: #165DFF; color: white; border: none; border-radius: 5px; cursor: pointer; }
                #status { color: #666; margin: 10px 0; }
            </style>
        </head>
        <body>
            <div class="container">
                <h1>Kokoro TTS 流式语音合成</h1>
                <textarea id="textInput">Kokoro 是一系列体积虽小但功能强大的 TTS 模型。该模型是经过短期训练的结果。</textarea>
                <div class="controls">
                    <label>音色：</label>
                    <select id="voiceSelect">
                        <option value="zf_001">女性音色 (zf_001)</option>
                        <option value="zm_010">男性音色 (zm_010)</option>
                    </select>
                    <label style="margin-left: 20px;">语速：</label>
                    <input type="range" id="speedRange" min="0.5" max="2.0" step="0.1" value="1.0">
                    <span id="speedValue">1.0</span>
                </div>
                <button onclick="startTTS()">开始合成并播放</button>
                <div id="status">就绪</div>
                <audio id="audioPlayer" controls style="width: 100%; margin-top: 20px;"></audio>
            </div>
            <script>
                // 显示当前语速值
                document.getElementById('speedRange').addEventListener('input', function() {
                    document.getElementById('speedValue').textContent = this.value;
                });
                
                // 开始TTS合成并播放
                async function startTTS() {
                    const text = document.getElementById('textInput').value;
                    const voice = document.getElementById('voiceSelect').value;
                    const speed = parseFloat(document.getElementById('speedRange').value);
                    const status = document.getElementById('status');
                    const audioPlayer = document.getElementById('audioPlayer');
                    
                    if (!text.trim()) {
                        alert('请输入文本');
                        return;
                    }
                    
                    status.textContent = '正在合成语音...';
                    
                    try {
                        // 调用流式接口
                        const response = await fetch('/api/tts/stream', {
                            method: 'POST',
                            headers: { 'Content-Type': 'application/json' },
                            body: JSON.stringify({ text, voice, speed })
                        });
                        
                        if (!response.ok) {
                            throw new Error(`HTTP错误：${response.status}`);
                        }
                        
                        // 读取流式响应并播放
                        const audioBlob = await response.blob();
                        const audioUrl = URL.createObjectURL(audioBlob);
                        audioPlayer.src = audioUrl;
                        audioPlayer.play();
                        status.textContent = '合成完成，正在播放...';
                        
                    } catch (error) {
                        status.textContent = `错误：${error.message}`;
                        console.error('TTS失败：', error);
                    }
                }
            </script>
        </body>
    </html>
    """

# 启动服务
if __name__ == "__main__":
    import uvicorn
    # 启动服务（host=0.0.0.0 允许外部访问，port=8000 端口）
    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)  # TTS模型不支持多进程，workers=1